from pymongo import MongoClient
from config.env import MongoConfig
from typing import Any, Dict
from datetime import datetime


class MongoDB:
    _client: MongoClient = None
    _db = None

    @classmethod
    def init(cls):
        cls._client = MongoClient(MongoConfig.mongo_url, maxPoolSize=100)
        cls._db = cls._client[MongoConfig.mongo_db]

    @classmethod
    def collection(cls, name: str):
        return cls._db[name]


class OcrMDB:

    _collection_name = "task"

    @classmethod
    def collection(cls):
        return MongoDB.collection(cls._collection_name)

    @classmethod
    def insert_one(cls, data):
        if "task_id" not in data:
            raise ValueError("task_id is required as primary key")
        time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        data["created_at"] = time
        data["updated_at"] = time
        return cls.collection().insert_one(data)

    @classmethod
    def update_one(cls, data):
        if "task_id" not in data:
            raise ValueError("task_id is required as primary key")
        data["updated_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        return cls.collection().update_one({"task_id": data["task_id"]}, {"$set": data})

    @classmethod
    def find_one(cls, query: Dict[str, Any]):
        return cls.collection().find_one(query)

    @classmethod
    def find(cls, query: Dict[str, Any]):
        return cls.collection().find(query)

    @classmethod
    def delete(cls, task_id: str):
        return cls.collection().delete_one({"task_id": task_id})


class TaskMDB(OcrMDB):
    _collection_name = "task"


class OcrTaskMDB(OcrMDB):
    _collection_name = "ocr-task"


class OcrResultMDB(OcrMDB):
    _collection_name = "ocr-result"


class LlmTaskMDB(OcrMDB):
    _collection_name = "llm-task"


class LlmResultMDB(OcrMDB):
    _collection_name = "llm-result"


class LlmChatMDB(OcrMDB):
    _collection_name = "llm-chat"
